from torch.utils.data import Dataset
import torch
import pathlib
import cv2
from typing import Tuple
import csv
import enum
import math
from torchvision.transforms import v2


class AutoLaparo(Dataset):
    """适用于SurgFormer的AutoLaparo数据集类
    
    按fps，每秒取第一帧

    返回的数据为: (clip_len, C, H, W) (1)
    """

    
    LABEL_MAPPING = {
        0: "Preparation", 1: "Dividing Ligament and Peritoneum", 2: "Dividing Uterine Vessels and Ligament", 
        3: "Transecting the Vagina", 4: "Specimen Removal", 5: "Suturing", 6: "Washing"
    }
    """标签id到标签名称的对应关系"""
    
    NUM_FRAMES_VIDEO = [
        6387, 3620, 2999, 2937, 3220, 3907, 1644, 4691, 5735, 5063, 4720,
        2915, 2596, 4738, 3652, 3611, 4677, 3545, 3412, 4831, 4325
    ]
    """从各个视频抽取的总帧数"""
    
    
    
    class DataStrategy(enum.Enum):
        """数据输出策略
        
        ONLINE时，当前/目标帧为最后的帧；OFFLINE时，当前/目标帧为中间的帧
        """
        ONLINE = "online"
        OFFLINE = "offline"

    def __init__(
        self, dataset_root: pathlib.Path, clip_len: int=16, data_strategy: DataStrategy=DataStrategy.ONLINE,
        resize: Tuple[int] | int=None
    ):
        super().__init__()
        
        self.clip_len = clip_len
        self.data_strategy = data_strategy
        
        self.root_path = dataset_root.absolute()
        self.label_path = dataset_root / "labels"
        self.frame_path = dataset_root / "frames"
        if isinstance(resize, int):
            self.resize = (resize, resize)
        else:
            self.resize = resize
        
        self.transforms = v2.Compose([
            v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
        
        # 存储所有视频所有帧的标签信息
        # video_id: 当前帧所在视频编号（1-index）
        # frame_id: 当前帧在所在视频中的编号（0-index），与帧图片文件对应
        # label_gt: 单个整数标记的ground truth label（按LABEL_MAPPING规则）
        self.labels = []
        # 读取所有标注标签
        for i in range(1, 22):
            anno_filepath = self.label_path / f"label_{i:02d}.txt"
            with open(anno_filepath, "r", encoding="utf-8") as fi:
                reader = csv.reader(fi, delimiter="\t")
                next(reader)
                for cnt, row in enumerate(reader):
                    if cnt >= AutoLaparo.NUM_FRAMES_VIDEO[i-1]:
                        break
                    self.labels.append({
                        'video_id': i,
                        'frame_id': cnt,
                        'label_gt': int(row[1])
                    })
    
    
    def __getitem__(self, target_idx):
        target_frame_info = self.labels[target_idx]
        img_sequence = []
        gt_label = torch.tensor(self.labels[target_idx]['label_gt'])
        
        if self.data_strategy == AutoLaparo.DataStrategy.ONLINE:
            # online策略的数据，即当前帧为最后一帧，往前多取clip_len-1帧，共clip_len帧
            # 所在视频的起始帧在整个数据集中的位置
            cur_video_start_idx = target_idx - target_frame_info['frame_id']

            # 当前要读取的全局编号
            cur_idx = target_idx
            while len(img_sequence) < self.clip_len:
                frame_info = self.labels[cur_idx]
                frame_path = self.frame_path / f"video{frame_info['video_id']:02d}_{frame_info['frame_id']:08d}.png"
                img = cv2.imread(frame_path.as_posix())
                if self.resize:
                    img = cv2.resize(img, tuple(reversed(self.resize)), interpolation=cv2.INTER_LINEAR)
                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                img = torch.tensor(img, dtype=torch.float32).permute(2, 0, 1).contiguous()
                img = self.transforms(img)
                img_sequence = [img] + img_sequence
                # 不超过所在视频，若前序帧不足，用视频第一帧填充
                if cur_idx > cur_video_start_idx:
                    cur_idx -= 1
        elif self.data_strategy == AutoLaparo.DataStrategy.OFFLINE:
            # offline策略的数据，即当前帧为中间帧，往前多取⌈(clip_len-1)/2⌉帧，往后多取⌊(clip_len-1)/2⌋帧，共clip_len帧
            left_len = math.floor((self.clip_len - 1) / 2)
            right_len = self.clip_len - 1 - left_len
            
            # 当前要读取的全局编号
            cur_idx = target_idx
            # 先取当前帧与前序帧
            while len(img_sequence) < (left_len + 1):
                frame_info = self.labels[cur_idx]
                frame_path = self.frame_path / f"video{frame_info['video_id']:02d}_{frame_info['frame_id']:08d}.png"
                img = cv2.imread(frame_path.as_posix())
                if self.resize:
                    img = cv2.resize(img, tuple(reversed(self.resize)), interpolation=cv2.INTER_LINEAR)
                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                img = torch.tensor(img, dtype=torch.float32).permute(2, 0, 1).contiguous()
                img = self.transforms(img)
                img_sequence = [img] + img_sequence
                # 不超过所在视频，若前序帧不足，用视频第一帧填充
                if cur_idx > 0 and self.labels[cur_idx]['video_id'] == self.labels[cur_idx-1]['video_id']:
                    cur_idx -= 1

            cur_idx = target_idx + 1 if target_idx < (len(self.labels) - 1) else target_idx
            # 再取后序帧
            while len(img_sequence) < self.clip_len:
                frame_info = self.labels[cur_idx]
                frame_path = self.frame_path / f"video{frame_info['video_id']:02d}_{frame_info['frame_id']:08d}.png"
                img = cv2.imread(frame_path.as_posix())
                if self.resize:
                    img = cv2.resize(img, tuple(reversed(self.resize)), interpolation=cv2.INTER_LINEAR)
                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                img = torch.tensor(img, dtype=torch.float32).permute(2, 0, 1).contiguous()
                img = self.transforms(img)
                img_sequence = img_sequence + [img]
                # 不超过所在视频，若后序帧不足，用视频最后一帧填充
                if cur_idx < (len(self.labels) - 1) and self.labels[cur_idx]['video_id'] == self.labels[cur_idx+1]['video_id']:
                    cur_idx += 1
        
        assert len(img_sequence) == self.clip_len
        
        img_sequence = torch.stack(img_sequence)
        
        return img_sequence, gt_label
    
    def __len__(self):
        return len(self.labels)